#ifndef TEST_HYPERCUBE_HYPERCUBE_TEST_HELPER
#define TEST_HYPERCUBE_HYPERCUBE_TEST_HELPER

#include "hypercube/hypercube_cryto_agent.hpp"
#include "hypercube/hypercube_hmm_engine.hpp"
#include "intel_itt_wrapper.hpp"
#include "log_helper.hpp"
#include "matrix_util.hpp"

#include <NTL/mat_ZZ.h>
#include <gtest/gtest.h>

static params<helib::BGV> my_params;

class MatMulTest : public testing::Test {
protected:
  void SetUp() override {
    ASSERT_PRINTF(!(m_ <= 0 || n_ <= 0 || k_ <= 0), "Invaild m, n or k\n");

    FIX_NTL_RAND_SEED(true); // fix random seed?
    RAND_NTL_MATRIX(A_, m_, k_, 8);
    RAND_NTL_MATRIX(B_, k_, n_, 8);

    // Get answer
    AB_ = A_ * B_;

    // Setup algorithm
    const auto eh_mnk = hypercube_hmme::get_expected_mnk(m_, n_, k_);
    const auto eh_cube = hypercube_hmme::get_hypercube(m_, n_, k_);
    hmmcc_ = new hypercube_hmmcc(my_params, eh_cube, cache_free_, minimal_m_);
    hmmcc_->report_context();

    hmme_ = new hypercube_hmme(hmm_status_);
    INTEL_ITT_RESUME;
    hmme_->register_engine(hmmcc_->context()->getEA(), eh_mnk);
    INTEL_ITT_PAUSE;

    // Get setup time
    setup_time = hmmcc_->setup_time();
    key_gen_time = hmmcc_->key_gen_time();
  }

  void TearDown() override {
    delete hmme_;
    delete hmmcc_;
  }

public:
  static void check(const NTL::mat_ZZ &ans, const NTL::mat_ZZ &res) {
    const long row = ans.NumRows();
    const long col = ans.NumCols();
    for (long i = 0; i < row; i++) {
      for (long j = 0; j < col; j++) {
        auto expect = ans[i][j];
        auto actual = res[i][j];
        ASSERT_EQ(expect, actual) << "diff at (" << i << ", " << j << ")";
      }
    }
  }

protected:
  // Actual dimension
  unsigned long m_ = 0;
  unsigned long n_ = 0;
  unsigned long k_ = 0;

  NTL::mat_ZZ A_;
  NTL::mat_ZZ B_;
  NTL::mat_ZZ AB_;

  MATRIX_TYPE type_ = MATRIX_TYPE::UNDEF_MATRIX_TYPE;
  hypercube_hmme *hmme_ = nullptr;
  hypercube_hmmcc *hmmcc_ = nullptr;
  hmm_status<shmm_engine> hmm_status_;

  double hmm_time = 0.0;
  double setup_time = 0.0;
  double key_gen_time = 0.0;

  bool cache_free_;
  bool minimal_m_;
};

#endif /* TEST_HYPERCUBE_HYPERCUBE_TEST_HELPER */
